package ark

import (
	"context"
	"os"

	"gitee.com/EagleMao/eino_n9e_tool/config"
	"gitee.com/EagleMao/eino_n9e_tool/pkg/logger"
	"github.com/cloudwego/eino-ext/components/model/ark"
	"go.uber.org/zap"
)

// CreateArkModel 创建ark大语言模型
func CreateArkModel(ctx context.Context, cfg config.LLMConfig) (*ark.ChatModel, error) {
	var apiKey = cfg.APIKey
	if os.Getenv("ARK_API_KEY") != "" {
		apiKey = os.Getenv("ARK_API_KEY")
	}
	if apiKey == "" || cfg.ModelName == "" {
		logger.Fatal("缺少ark大语言模型配置")
		return nil, nil
	}
	// export ARK_API_KEY=7d7e766b-5659-4020-ac9d-74de7af3ff4d
	model, err := ark.NewChatModel(
		ctx,
		&ark.ChatModelConfig{
			APIKey: apiKey,
			Model:  cfg.ModelName,
		},
	)

	if err != nil {
		logger.Fatal("创建ark大语言模型失败", zap.Error(err))
		return nil, err
	}

	logger.Info("创建ark大语言模型成功")

	return model, nil
}
